use fxhash::FxHashMap;



use crate::gpu::program::assets::ShaderLib;

use super::StringUtil;

/// shader 预处理器
pub struct PreProcessor;

impl PreProcessor {
    pub fn parse(
        shader_lib: &ShaderLib,
        code: &str,
        define_value: &mut FxHashMap<String, String>,
    ) -> String {
        //
        let code = Self::filter_comment(code);
        let mut ctx = PreProcessorContext::default();
        let code = Self::parse_preprocess(shader_lib, &mut ctx, &code, define_value);
        Self::parse_auto_binding_for_group_x(&code, 1)
    }

    pub fn parse_compute_shader(
        shader_lib: &ShaderLib,
        code: &str,
        define_value: &mut FxHashMap<String, String>,
    ) -> String {
        let code = Self::filter_comment(code);
        let mut ctx = PreProcessorContext::default();
        Self::parse_preprocess(shader_lib, &mut ctx, &code, define_value)
    }

    pub fn parse_preprocess(
        shader_lib: &ShaderLib,
        context: &mut PreProcessorContext,
        code: &str,
        define_value: &mut FxHashMap<String, String>,
    ) -> String {
        if let Some(index) = code.find('#') {
            let header = &code[0..index];
            let last_icon_pos = code.rfind('#').unwrap();
            let temp_code = &code[last_icon_pos..];
            
            let end_index = temp_code.find('\n').unwrap_or(temp_code.len()) + last_icon_pos;
            let code_block = &code[index..end_index];

            let tail = &code[end_index..];
            let content =
                Self::parse_preprocess_command(shader_lib, context, code_block, define_value);

            format!("{header}{content}{tail}")
        } else {
            code.to_owned()
        }
    }

    pub fn parse_auto_binding_for_group_x(code: &str, n_group: u32) -> String {
        let mut offset = 0;
        let mut result = "".to_owned();
        let mut group = FxHashMap::default();
        let len = code.len();
        while offset < len {
            if let Some(n_left_index) = StringUtil::str_find_at(code, "@group(", offset) {
                let n_right_index = StringUtil::str_find_at(code, ')', n_left_index).unwrap();
                let group_id = code[(n_left_index + 7)..n_right_index]
                    .parse::<u32>()
                    .unwrap();
                let n_left_index =
                    StringUtil::str_find_at(code, "@binding(", n_right_index).unwrap();
                let n_right_index = StringUtil::str_find_at(code, ")", n_left_index).unwrap();

                result.push_str(&code[offset..n_left_index]);

                if group_id == n_group {
                    match group.entry(group_id) {
                        std::collections::hash_map::Entry::Occupied(mut group) => {
                            let id = group.get() + 1;
                            result.push_str(&format!("@binding({id})"));
                            *group.get_mut() = id;
                        }
                        std::collections::hash_map::Entry::Vacant(group) => {
                            result.push_str("@binding(0)");
                            group.insert(0);
                        }
                    }
                } else {
                    result.push_str(&code[n_left_index..n_right_index + 1]);
                }
                offset = n_right_index + 1;
            } else {
                result.push_str(&code[offset..]);
                break;
            }
        }
        result
    }

    pub fn parse_preprocess_command(
        shader_lib: &ShaderLib,
        context: &mut PreProcessorContext,
        code: &str,
        define_value: &mut FxHashMap<String, String>,
    ) -> String {
        let mut result = "".to_owned();
        let lines = code.split('\n');

        let mut stacks = vec![false];
        let mut stack_else_ifs = vec![false];

        for line in lines {
            let mut skip = stacks[stacks.len() - 1];
            let pos = line.trim().find('#').unwrap_or(usize::MAX) ;
            {
                if pos != 0 {
                    if !skip {
                        result.push_str(&format!("{}\n", line));
                    }
                    continue;
                }
            }

            let command = line.trim();
            if command.contains("#if") {
                if skip && stacks.len() > 1 {
                    stacks.push(skip);
                    continue;
                }
                let condition = command[3..].trim();
                skip = !Self::parse_condition(condition, define_value);
                stacks.push(skip);
                stack_else_ifs.push(!skip);
                continue;
            } else if command.contains("#elseif")
                // || command.find("#else").is_some()
                // || command.contains(" if")
            {
                let skip_else_if = stack_else_ifs[stack_else_ifs.len() - 1];
                if skip_else_if {
                    stacks.pop();
                    skip = true;
                    stacks.push(skip);
                    continue;
                }
                stacks.pop();
                skip = *stacks.last().unwrap();
                if skip && stacks.len() > 1 {
                    stacks.push(skip);
                    continue;
                }
                // println!("Command = {}",command);
                let condition = command[(command.find("if").unwrap() + 2)..].trim();

                if condition.is_empty() {
                    println!(
                        "preprocess command error, conditions missing: {:?}",
                        command
                    );
                    continue;
                }

                skip = !Self::parse_condition(condition, define_value);
                stacks.push(skip);
                stack_else_ifs.push(!skip);
                continue;
            } else if command.contains("#else") {
                let skip_else_if = stack_else_ifs[stack_else_ifs.len() - 1];
                stacks.pop();
                if stack_else_ifs.iter().any(|cond|*cond == true) {
                    
                    skip = true;
                    stacks.push(skip);
                    continue;
                }else {
                    skip = false;
                    stacks.push(skip);
                    continue;
                }
                
                // if skip && stacks.len() > 1 && stacks.last().is_some() {
                //     stacks.push(skip);
                // } else {
                //     stacks.push(!skip);
                // }
            } else if command.find("#endif").is_some() {
                stacks.pop();
                stack_else_ifs.pop();
                continue;
            } else if command.find("#include").is_some() {
                let mut include_name = "".to_owned();
                if let Some(chr) = command.chars().last() {
                    if chr == '>' {
                        include_name = Self::extract(command, "<", ">");
                    } else {
                        include_name = Self::extract(command, &chr.to_string(), &chr.to_string())
                    }
                }

                if !context.include_map.contains_key(&include_name) {
                    let code = shader_lib.get_shader(&include_name);
                    context.include_map.insert(include_name.clone(), true);
                    if code.is_empty() {
                        panic!("{command} error: {include_name} not found;")
                    }

                    let code = Self::filter_comment(&code);
                    let code = Self::parse_preprocess(shader_lib, context, &code, define_value);
                    result.push_str(&format!("{code}\r\n"));
                }
                continue;
            } else if let Some(define_pos) = command.find("#define ") {
                let expr = command[(define_pos + 8)..].trim();

                let mut name = expr.to_owned();
                let mut value = "".to_owned();

                if let Some(index) = expr.find(" ") {
                    name = expr[0..index].to_owned();
                    value = expr[(index + 1)..].to_owned();
                }

                define_value.insert(name, value);
                continue;
            } else {
                panic!("nonsupport: {command}")
            }
        }

        result
    }

    pub fn extract(str: &str, left: &str, right: &str) -> String {
        let index_l = str.find(left).unwrap() + left.len();
        let temp_str = &str[index_l..];
        let index_r = temp_str.find(right).unwrap() + index_l;
        str[index_l..index_r].to_owned()
    }

    pub fn parse_condition(condition: &str, define_value: &FxHashMap<String, String>) -> bool {
        if let Some(value) = define_value.get(condition) {
            value == "true"
        } else {
            false
        }
    }

    pub fn filter_comment(code: &str) -> String {
        let mut result = "".to_owned();
        let mut find_single_comment = true;
        let mut find_multi_comment = true;

        let mut offset = 0;
        let len = code.len();
        while offset < len {
            let temp_code = &code[offset..];
            let mut index1 = if find_single_comment {
                temp_code.find("//")
            } else {
                None
            };

            let mut index2 = if find_multi_comment {
                temp_code.find("/*")
            } else {
                None
            };

            if index1.is_none() && index2.is_none() {
                result.push_str(temp_code);
                break;
            }

            find_single_comment = index1.is_some();
            find_multi_comment = index2.is_some();

            if index1.is_some() && index2.is_some() {
                if index1.unwrap() < index2.unwrap() {
                    index2 = None;
                } else {
                    index1 = None;
                }
            }

            if let Some(index1) = index1 {
                let temp_code1 = &temp_code[index1..];
                let code_offset = offset + index1;
                let index2 = temp_code1.find("\n");
                result.push_str(&code[offset..code_offset]);

                offset = if let Some(index2) = index2 {
                    code_offset + index2
                } else {
                    code.len()
                }
            } else if let Some(index2) = index2 {
                let new_code = &temp_code[index2..];
                let code_offset = offset + index2;
                let index1 = new_code.find("*/").unwrap();

                result.push_str(&code[offset..code_offset]);

                offset = code_offset + index1 + 2;
            }
        }

        result
    }
}

#[derive(Default)]
pub struct PreProcessorContext {
    pub include_map: FxHashMap<String, bool>,
}
